/**************************************************************************** 
 * Meta，一个轻量级的静态分区Hypervisor 
 *
 * 版权(c) , 2022-
 *
 * 作者:
 *      Tupelo Shen <shenwanjiang2013@163.com>
 *
 * Meta是一个自由软件，你可以在遵循GNU GPLv2协议的条件下，重新修改并发布它。
 *
 ***************************************************************************/

#include <arch/psci.h>
#include <arch/smc.h>
#include <arch/gic.h>
#include <meta.h>
#include <fences.h>
#include <cpu.h>
#include <vm.h>

/****************************************************************************
 ****                           私有函数接口                             ****
 ****************************************************************************/
static void psci_save_state(enum wakeup_reason wakeup_reason);
static void psci_restore_state(void);
static void psci_wake_from_off(void);
static void psci_wake_from_powerdown(void);
static void psci_wake_from_idle(void);

/****************************************************************************
 ****                           变量声明                                 ****
 ****************************************************************************/
void (*psci_wake_handlers[PSCI_WAKEUP_NUM])(void) = {
    [PSCI_WAKEUP_CPU_OFF]   = psci_wake_from_off,
    [PSCI_WAKEUP_POWERDOWN] = psci_wake_from_powerdown,
    [PSCI_WAKEUP_IDLE]      = psci_wake_from_idle,
};

enum {PSCI_MSG_ON};

extern uint8_t  root_l1_flat_pt;
extern void     psci_boot_entry(unsigned long x0);
/****************************************************************************
 ****                           私有函数定义                             ****
 ****************************************************************************/

/**
 * 功能描述: 保存当前运行状态
 *           (1) 保存hyp自身MMU相关寄存器
 *           (2) 保存虚拟化相关的MMU寄存器
 *           (3) 保存root平面映射页表
 *           (4) 保存唤醒之前的状态
 *           (5) 刷新cache，保证前面的保存成功写入内存
 *           (6) 保存GICC状态（也就是`CPU interface`相关寄存器）
 * 入口参数: wakeup_reason:     唤醒之前的状态
 * 出口参数: void
 * 返 回 值: void
 * 修改日期: 2022-09-21
 */
static void psci_save_state(enum wakeup_reason wakeup_reason)
{
    cpu.arch.psci_off_state.tcr_el2     = MRS(TCR_EL2);
    cpu.arch.psci_off_state.ttbr0_el2   = MRS(TTBR0_EL2);
    cpu.arch.psci_off_state.mair_el2    = MRS(MAIR_EL2);
    cpu.arch.psci_off_state.cptr_el2    = MRS(CPTR_EL2);
    cpu.arch.psci_off_state.hcr_el2     = MRS(HCR_EL2);
    cpu.arch.psci_off_state.vmpidr_el2  = MRS(VMPIDR_EL2);
    cpu.arch.psci_off_state.vtcr_el2    = MRS(VTCR_EL2);
    cpu.arch.psci_off_state.vttbr_el2   = MRS(VTTBR_EL2);

    /* 将root平面映射页表存储到cpu数据结构中，作为唤醒时的页表 */
    mem_translate(&cpu.as,
                  (vaddr_t)&root_l1_flat_pt,
                  &cpu.arch.psci_off_state.flat_map);
    cpu.arch.psci_off_state.wakeup_reason = wakeup_reason;

    /**
     * Although the real PSCI implementation is responsible for managing cache
     * state, make sure the saved state is in memory as we'll use this on wake
     * up before enabling cache to restore basic processor state. 
     */
    cache_flush_range((vaddr_t)&cpu.arch.psci_off_state,
                      sizeof(cpu.arch.psci_off_state));

    /* 保存中断信息 */
    gicc_save_state(&cpu.arch.psci_off_state.gicc_state);
}

/**
 * 功能描述: 恢复之前的状态
 * 入口参数: void
 * 出口参数: void
 * 返 回 值: void
 * 修改日期: 2022-09-21
 */
static void psci_restore_state(void)
{

    /* 要恢复的状态，大部分已经在汇编函数psci_boot_entry实现。
     * 此处，恢复中断状态信息
     */
    gicc_restore_state(&cpu.arch.psci_off_state.gicc_state);
}

/**
 * 功能描述: 从CPU_OFF状态唤醒
 * 入口参数: void
 * 出口参数: void
 * 返 回 值: void
 * 修改日期: 2022-09-21
 */
static void psci_wake_from_off(void)
{
    /* 如果当前核没有虚拟核，则无需处理 */
    if(cpu.vcpu == NULL){
        return;
    }

    /* 否则，更新vcpu.psci_ctx */
    spin_lock(&cpu.vcpu->arch.psci_ctx.lock);
    if(cpu.vcpu->arch.psci_ctx.state == ON_PENDING){
        vcpu_arch_reset(cpu.vcpu, cpu.vcpu->arch.psci_ctx.entrypoint);
        cpu.vcpu->arch.psci_ctx.state = ON;
        cpu.vcpu->regs->x[0] = cpu.vcpu->arch.psci_ctx.context_id;
    }
    spin_unlock(&cpu.vcpu->arch.psci_ctx.lock);
}

/**
 * 功能描述: 从POWER_DOWN状态唤醒
 * 入口参数: void
 * 出口参数: void
 * 返 回 值: void
 * 修改日期: 2022-09-21
 */
static void psci_wake_from_powerdown(void)
{

    if(cpu.vcpu == NULL){
        ERROR("cpu woke up but theres no vcpu to run");
    }

    vcpu_arch_reset(cpu.vcpu, cpu.vcpu->arch.psci_ctx.entrypoint);
    vcpu_writereg(cpu.vcpu, 0, cpu.vcpu->arch.psci_ctx.context_id);
    vcpu_run(cpu.vcpu);
}

/**
 * 功能描述: 从IDLE状态唤醒
 * 入口参数: void
 * 出口参数: void
 * 返 回 值: void
 * 修改日期: 2022-09-21
 */
static void psci_wake_from_idle(void)
{
    cpu_idle_wakeup();
}
// int32_t psci_standby(){
//     /* only apply request to core level */
//     uint32_t pwr_state_aux = PSCI_POWER_STATE_LVL_0 | PSCI_STATE_TYPE_STANDBY;

//     return psci_cpu_suspend(pwr_state_aux, 0, 0);
// }

/****************************************************************************
 ****                           公共函数定义                             ****
 ****************************************************************************/
/**
 * 功能描述: cpu核被唤醒后的处理程序
 * 入口参数: handler_id:  处理的句柄
 * 出口参数: void
 * 返 回 值: void
 * 修改日期: 2022-09-21
 */
void psci_wake(uint32_t handler_id)
{    

    psci_restore_state();

    if(handler_id < PSCI_WAKEUP_NUM){
        psci_wake_handlers[handler_id]();
    } else {
        ERROR("unkown reason for cpu wake up");
    }
}

/**
 * 功能描述: 将cpuid根据配置的cluster信息转换成mpidr值
 * 入口参数: reason:    cpu核掉电的原因
 * 出口参数: void
 * 返 回 值: uint64_t：返回该核对应的mpidr值
 * 修改日期: 2022-09-21
 */
int32_t psci_power_down(enum wakeup_reason reason)
{
    /* 电源状态 TODO: 此处应该有问题 */
    uint32_t pwr_state_aux = PSCI_POWER_STATE_LVL_0 | PSCI_STATE_TYPE_POWERDOWN;

    psci_save_state(reason);
    paddr_t cntxt_paddr;
    paddr_t psci_wakeup_addr;
    mem_translate(&cpu.as, (vaddr_t)&cpu.arch.psci_off_state, &cntxt_paddr);
    mem_translate(&cpu.as, (vaddr_t)&psci_boot_entry, &psci_wakeup_addr);

    return psci_cpu_suspend(pwr_state_aux, psci_wakeup_addr, cntxt_paddr);
}

/*********************** SMC PSCI interface  ********************************/
/**
 * 功能描述: 通过SMC调用，启动一个核
 * 入口参数: target_cpu:    目标核
 *           entrypoint:    入口地址
 *           context_id:    
 * 出口参数: void
 * 返 回 值: uint64_t：返回结果
 * 修改日期: 2022-09-21
 */
int32_t psci_cpu_on(unsigned long target_cpu, unsigned long entrypoint, 
                    unsigned long context_id)
{
    return smc_call(PSCI_CPU_ON_AARCH64, target_cpu, entrypoint, context_id,
                    NULL);
}

/**
 * 功能描述: 通过SMC调用，关闭一个核
 *           TODO: 暂时未使用
 * 入口参数: void    
 * 出口参数: void
 * 返 回 值: uint64_t：返回结果
 * 修改日期: 2022-09-21
 */
int32_t psci_cpu_off(void)
{   
    return smc_call(PSCI_CPU_OFF, 0, 0, 0, NULL);
}

/**
 * 功能描述: 通过SMC调用，挂起一个核
 * 入口参数: power_state:   目标核
 *           entrypoint:    入口地址
 *           context_id:   
 * 出口参数: void
 * 返 回 值: uint64_t：返回结果
 * 修改日期: 2022-09-21
 */
int32_t psci_cpu_suspend(uint32_t power_state, unsigned long entrypoint, 
                    unsigned long context_id)
{

    return smc_call(PSCI_CPU_SUSPEND_AARCH64,
                    power_state,
                    entrypoint, 
                    context_id,
                    NULL);
}

/**
 * 功能描述: 通过SMC调用，获取psci版本
 * 入口参数: void   
 * 出口参数: void
 * 返 回 值: uint64_t：返回结果
 * 修改日期: 2022-09-21
 */
int32_t psci_version(void)
{
    return smc_call(PSCI_VERSION, 0, 0, 0, NULL);
}

/**
 * 功能描述: guest os的psci_cpu_off请求处理
 * 入口参数: void   
 * 出口参数: void
 * 返 回 值: uint64_t：返回结果
 * 修改日期: 2022-09-21
 */
static int32_t psci_cpu_off_handler(void)
{
    /**
     *  Right now we only support one vcpu por cpu, so passthrough the request
     *  directly to the monitor psci implementation. Later another vcpu, will
     *  call cpu_on on this vcpu. 
     */

    spin_lock(&cpu.vcpu->arch.psci_ctx.lock);
    cpu.vcpu->arch.psci_ctx.state = OFF;
    spin_unlock(&cpu.vcpu->arch.psci_ctx.lock);

    cpu_idle();

    spin_lock(&cpu.vcpu->arch.psci_ctx.lock);
    cpu.vcpu->arch.psci_ctx.state = ON;
    spin_unlock(&cpu.vcpu->arch.psci_ctx.lock);

    return PSCI_E_DENIED;
}

/**
 * 功能描述: guest os的psci_cpu_suspend请求处理
 * 入口参数: void   
 * 出口参数: void
 * 返 回 值: uint64_t：返回结果
 * 修改日期: 2022-09-21
 */
static int32_t psci_cpu_suspend_handler(uint32_t        power_state,
                                        unsigned long   entrypoint,
                                        unsigned long   context_id)
{
    /**
     * !! Ignoring the rest of the requested  powerstate for now. 
     * This might be a problem howwver since powerlevel and stateid are 
     * implementation defined.
     */ 
    uint32_t state_type = power_state & PSCI_STATE_TYPE_BIT;
    int32_t ret;

    if(state_type){
        //PSCI_STATE_TYPE_POWERDOWN:
        spin_lock(&cpu.vcpu->arch.psci_ctx.lock);
        cpu.vcpu->arch.psci_ctx.entrypoint = entrypoint;
        cpu.vcpu->arch.psci_ctx.context_id = context_id;
        spin_unlock(&cpu.vcpu->arch.psci_ctx.lock);
        ret = psci_power_down(PSCI_WAKEUP_POWERDOWN);
    } else {
        //PSCI_STATE_TYPE_STANDBY:
        /**
         *  TODO: ideally we would emmit a standby request to PSCI 
         * (currently, ATF), but when we do, we do not wake up on interrupts
         * on the current development target zcu104.
         * We should understand why. To circunvent this, we directly emmit a
         * wfi 
         */
        //ret = psci_standby();
        asm volatile("wfi\n\r");
        ret = PSCI_E_SUCCESS;
    }

    return ret;
}

static void psci_cpumsg_handler(uint32_t event, uint64_t data){

    switch(event){
        case PSCI_MSG_ON:
            psci_wake_from_off();
        break;
    }
}

/**
 * 功能描述: guest os的psci_cpu_suspend请求处理
 * 入口参数: void   
 * 出口参数: void
 * 返 回 值: uint64_t：返回结果
 * 修改日期: 2022-09-21
 */
CPU_MSG_HANDLER(psci_cpumsg_handler, PSCI_CPUSMG_ID);
static int32_t psci_cpu_on_handler( unsigned long   target_cpu,
                                    unsigned long   entrypoint,
                                    unsigned long   context_id)
{
    int32_t ret;
    struct vm* vm = cpu.vcpu->vm;
    struct vcpu* target_vcpu = vm_get_vcpu_by_mpidr(vm, target_cpu);

    if (target_vcpu != NULL){

        bool already_on = true;
        spin_lock(&cpu.vcpu->arch.psci_ctx.lock);
        if(target_vcpu->arch.psci_ctx.state == OFF){
            target_vcpu->arch.psci_ctx.state = ON_PENDING;
            target_vcpu->arch.psci_ctx.entrypoint = entrypoint;
            target_vcpu->arch.psci_ctx.context_id = context_id;
            fence_sync_write();
            already_on = false;
        } 
        spin_unlock(&cpu.vcpu->arch.psci_ctx.lock);

        if(already_on){
            return PSCI_E_ALREADY_ON;
        }

        cpuid_t pcpuid = vm_translate_to_pcpuid(vm, target_vcpu->id);
        if (pcpuid == INVALID_CPUID) {
            ret = PSCI_E_INVALID_PARAMS;
        } else {
            struct cpu_msg msg = {PSCI_CPUSMG_ID, PSCI_MSG_ON};
            cpu_send_msg(pcpuid, &msg);
            ret = PSCI_E_SUCCESS;
        }
        // cpu_msg_t msg = {
        //     PSCI_CPUSMG_ID,
        //     PSCI_MSG_ON
        // };
        // cpu_send_msg(pcpuid, &msg);

        ret = PSCI_E_SUCCESS;

    } else {
        ret = PSCI_E_INVALID_PARAMS;
    }

    return ret;
}

/**
 * 功能描述: guest os的psci_affinity_info请求处理
 * 入口参数: void   
 * 出口参数: void
 * 返 回 值: uint64_t：返回结果
 * 修改日期: 2022-09-21
 */
static int32_t psci_affinity_info_handler(  unsigned long target_affinity, 
                                            uint32_t lowest_affinity_level)
{
    /* return ON, if at least one core in the affinity instance: has been 
    enabled with a call to CPU_ON, and that core has not called CPU_OFF */

    /* return off if all of the cores in the affinity instance have called 
    CPU_OFF and each of these calls has been processed by the PSCI 
    implementation. */

    /*  return ON_PENDING if at least one core in the affinity instance is in 
    the ON_PENDING state */

    /**
     * TODO
     */

    return 0;
}

/**
 * 功能描述: guest os的psci_features请求处理
 * 入口参数: void   
 * 出口参数: void
 * 返 回 值: uint64_t：返回结果
 * 修改日期: 2022-09-21
 */
static int32_t psci_features_handler(uint32_t feature_id)
{

    int32_t ret = PSCI_E_NOT_SUPPORTED;

    switch (feature_id) {
        case PSCI_VERSION:
        case PSCI_CPU_OFF:
        case PSCI_CPU_SUSPEND_AARCH64:
        case PSCI_CPU_ON_AARCH64:
        case PSCI_AFFINITY_INFO_AARCH32:
        case PSCI_FEATURES:
            ret = PSCI_E_SUCCESS;
            break;
    }    

    return ret;
}

/**
 * 功能描述: 通过SMC调用，获取目标核的亲和力信息
 *           TODO: 暂时未使用
 * 入口参数: target_affinity:   目标亲和力
 *           lowest_affinity_level: 最低亲和力等级   
 * 出口参数: void
 * 返 回 值: uint64_t：返回结果
 * 修改日期: 2022-09-21
 */
int32_t psci_affinity_info(unsigned long target_affinity, 
                           uint32_t lowest_affinity_level)
{
    return smc_call(PSCI_AFFINITY_INFO_AARCH64, target_affinity, 
                                            lowest_affinity_level, 0, NULL);
}

/**
 * 功能描述: psci规范中对smc异常的处理函数
 * 入口参数: smc_fid,       smc功能ID
 *           x1,            参数1
 *           x2,            参数2
 *           x3,            参数3
 * 出口参数: void
 * 返 回 值: int64_t,  psci调用的返回值
 * 修改日期: 2022-09-21
 */
int32_t psci_smc_handler(uint32_t smc_fid,
                        unsigned long x1,
                        unsigned long x2,
                        unsigned long x3)
{
   int32_t ret = PSCI_E_NOT_SUPPORTED;

    switch (smc_fid) {
        case PSCI_VERSION:
            ret = psci_version();
            break;

        case PSCI_CPU_OFF:
            ret = psci_cpu_off_handler();
            break;

        case PSCI_CPU_SUSPEND_AARCH64:
            ret = psci_cpu_suspend_handler(x1, x2, x3);
            break;

        case PSCI_CPU_ON_AARCH64:
            ret = psci_cpu_on_handler(x1, x2, x3);
            break;

        case PSCI_AFFINITY_INFO_AARCH64:
            ret = psci_affinity_info_handler(x1, x2);
            break;

        case PSCI_FEATURES:
            ret = psci_features_handler(x1);
            break;

        case PSCI_MIG_INFO_TYPE:
            ret = PSCI_TOS_NOT_PRESENT_MP;
            break;

        default:
            INFO("unkown psci smc_fid 0x%lx", smc_fid);
   }

    return ret;
}